"""
数据库迁移脚本
为sessions表添加category字段
"""

from sqlalchemy import create_engine, MetaData, Table, Column, String, text
from src.config.database import Base
from src.models.session import Session
import os
import sys

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.config.settings import settings

def upgrade():
    """添加category字段到sessions表"""
    engine = create_engine(getattr(settings, 'DATABASE_URL'))
    metadata = MetaData()
    
    # 反射现有的表结构
    metadata.reflect(bind=engine)
    
    # 获取sessions表
    if 'sessions' in metadata.tables:
        sessions_table = metadata.tables['sessions']
        
        # 检查category列是否已存在
        if 'category' not in sessions_table.c:
            # 添加category列
            with engine.connect() as conn:
                # 使用ALTER TABLE语句添加列
                try:
                    conn.execute(
                        text("ALTER TABLE sessions ADD COLUMN category VARCHAR(20) DEFAULT 'default'")
                    )
                    conn.commit()
                    print("成功添加category字段到sessions表")
                except Exception as e:
                    # 如果是SQLite，可能需要特殊处理
                    if "sqlite" in getattr(settings, 'DATABASE_URL').lower():
                        print("SQLite数据库需要特殊处理，建议重新创建表")
                        print("错误信息:", str(e))
                    else:
                        raise e
        else:
            print("category字段已存在")
    else:
        print("sessions表不存在")

def downgrade():
    """移除category字段（此操作在SQLite中可能不可行）"""
    print("注意：SQLite不支持直接删除列操作")
    print("如需回滚，请手动处理数据库表结构")

if __name__ == "__main__":
    upgrade()